package edu.vub.at.kryo;

import static com.esotericsoftware.kryo.Kryo.NULL;

import java.io.ObjectStreamException;
import java.lang.reflect.Array;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.Serializer;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;

public class ExtendedObjectSerializers {

	static public class ObjectArraySerializer extends Serializer<Object[]> {
		private boolean elementsAreSameType;
		private boolean elementsCanBeNull = true;

		{
			setAcceptsNull(true);
		}

		public void write (Kryo kryo, Output output, Object[] object) {
			if (object == null) {
				output.writeByte(NULL);
				return;
			}
			try{
			output.writeInt(object.length + 1, true);
			Class elementClass = object.getClass().getComponentType();
			if (elementsAreSameType || Modifier.isFinal(elementClass.getModifiers())) {
				Serializer elementSerializer = kryo.getSerializer(elementClass);
				for (int i = 0, n = object.length; i < n; i++) {
					if (elementsCanBeNull)
						kryo.writeObjectOrNull(output, ATKryo.writeReplace(object[i]), elementSerializer);
					else
						kryo.writeObject(output, ATKryo.writeReplace(object[i]), elementSerializer);
				}
			} else {
				for (int i = 0, n = object.length; i < n; i++)
					kryo.writeClassAndObject(output, ATKryo.writeReplace(object[i]));
			}
		}catch(ObjectStreamException e){
			e.printStackTrace();
		}
		}

		public Object[] read (Kryo kryo, Input input, Class<Object[]> type) {
			int length = input.readInt(true);
			if (length == NULL) return null;
			Object[] object = (Object[])Array.newInstance(type.getComponentType(), length - 1);
			kryo.reference(object);
			Class elementClass = object.getClass().getComponentType();
			
			try{
			if (elementsAreSameType || Modifier.isFinal(elementClass.getModifiers())) {
				Serializer elementSerializer = kryo.getSerializer(elementClass);
				for (int i = 0, n = object.length; i < n; i++) {
					if (elementsCanBeNull)
						object[i] = ATKryo.readResolve(kryo.readObjectOrNull(input, elementClass, elementSerializer));
					else
						object[i] = ATKryo.readResolve(kryo.readObject(input, elementClass, elementSerializer));
				}
			} else {
				for (int i = 0, n = object.length; i < n; i++)
					object[i] = ATKryo.readResolve(kryo.readClassAndObject(input));
			}
		}catch(ObjectStreamException e){
			e.printStackTrace();
		}
		return object;

		}

		public Object[] copy (Kryo kryo, Object[] original) {
			Object[] copy = (Object[])Array.newInstance(original.getClass().getComponentType(), original.length);
			System.arraycopy(original, 0, copy, 0, copy.length);
			return copy;
		}

		/** @param elementsCanBeNull False if all elements are not null. This saves 1 byte per element if the array type is final or
		 *           elementsAreSameClassAsType is true. True if it is not known (default). */
		public void setElementsCanBeNull (boolean elementsCanBeNull) {
			this.elementsCanBeNull = elementsCanBeNull;
		}

		/** @param elementsAreSameType True if all elements are the same type as the array (ie they don't extend the array type). This
		 *           saves 1 byte per element if the array type is not final. Set to false if the array type is final or elements
		 *           extend the array type (default). */
		public void setElementsAreSameType (boolean elementsAreSameType) {
			this.elementsAreSameType = elementsAreSameType;
		}
	}
	
	public class CollectionSerializer extends Serializer<Collection> {
		private boolean elementsCanBeNull = true;
		private Serializer serializer;
		private Class elementClass;
		private Class genericType;

		public CollectionSerializer () {
		}

		/** @see #setElementClass(Class, Serializer) */
		public CollectionSerializer (Class elementClass, Serializer serializer) {
			setElementClass(elementClass, serializer);
		}

		/** @see #setElementClass(Class, Serializer)
		 * @see #setElementsCanBeNull(boolean) */
		public CollectionSerializer (Class elementClass, Serializer serializer, boolean elementsCanBeNull) {
			setElementClass(elementClass, serializer);
			this.elementsCanBeNull = elementsCanBeNull;
		}

		/** @param elementsCanBeNull False if all elements are not null. This saves 1 byte per element if elementClass is set. True if it
		 *           is not known (default). */
		public void setElementsCanBeNull (boolean elementsCanBeNull) {
			this.elementsCanBeNull = elementsCanBeNull;
		}

		/** @param elementClass The concrete class of each element. This saves 1-2 bytes per element. Set to null if the class is not
		 *           known or varies per element (default).
		 * @param serializer The serializer to use for each element. */
		public void setElementClass (Class elementClass, Serializer serializer) {
			this.elementClass = elementClass;
			this.serializer = serializer;
		}

		public void setGenerics (Kryo kryo, Class[] generics) {
			if (kryo.isFinal(generics[0])) genericType = generics[0];
		}

		public void write (Kryo kryo, Output output, Collection collection) {
			int length = collection.size();
			output.writeInt(length, true);
			Serializer serializer = this.serializer;
			if (genericType != null) {
				if (serializer == null) serializer = kryo.getSerializer(genericType);
				genericType = null;
			}
			try{
			if (serializer != null) {
				if (elementsCanBeNull) {
					for (Object element : collection)
						kryo.writeObjectOrNull(output, ATKryo.writeReplace(element), serializer);
				} else {
					for (Object element : collection)
						kryo.writeObject(output, ATKryo.writeReplace(element), serializer);
				}
			} else {
				for (Object element : collection)
					kryo.writeClassAndObject(output, ATKryo.writeReplace(element));
			}
			}catch(ObjectStreamException e){
				e.printStackTrace();
			}
		}

		/** Used by {@link #read(Kryo, Input, Class)} to create the new object. This can be overridden to customize object creation, eg
		 * to call a constructor with arguments. The default implementation uses {@link Kryo#newInstance(Class)}. */
		protected Collection create (Kryo kryo, Input input, Class<Collection> type) {
			return kryo.newInstance(type);
		}

		public Collection read (Kryo kryo, Input input, Class<Collection> type) {
			Collection collection = create(kryo, input, type);
			kryo.reference(collection);
			int length = input.readInt(true);
			if (collection instanceof ArrayList) ((ArrayList)collection).ensureCapacity(length);
			Class elementClass = this.elementClass;
			Serializer serializer = this.serializer;
			if (genericType != null) {
				if (serializer == null) {
					elementClass = genericType;
					serializer = kryo.getSerializer(genericType);
				}
				genericType = null;
			}
		try{
			if (serializer != null) {
				if (elementsCanBeNull) {
					for (int i = 0; i < length; i++)
						collection.add(ATKryo.readResolve(kryo.readObjectOrNull(input, elementClass, serializer)));
				} else {
					for (int i = 0; i < length; i++)
						collection.add(ATKryo.readResolve(kryo.readObject(input, elementClass, serializer)));
				}
			} else {
				for (int i = 0; i < length; i++)
					collection.add(ATKryo.readResolve(kryo.readClassAndObject(input)));
			}
		}catch(ObjectStreamException e){
			e.printStackTrace();
		}
			return collection;
		}

		/** Used by {@link #copy(Kryo, Collection)} to create the new object. This can be overridden to customize object creation, eg to
		 * call a constructor with arguments. The default implementation uses {@link Kryo#newInstance(Class)}. */
		protected Collection createCopy (Kryo kryo, Collection original) {
			return kryo.newInstance(original.getClass());
		}

		public Collection copy (Kryo kryo, Collection original) {
			Collection copy = createCopy(kryo, original);
			kryo.reference(copy);
			for (Object element : original)
				copy.add(kryo.copy(element));
			return copy;
		}
	}
	
	public class MapSerializer extends Serializer<Map> {
		private Class keyClass, valueClass;
		private Serializer keySerializer, valueSerializer;
		private boolean keysCanBeNull = true, valuesCanBeNull = true;
		private Class keyGenericType, valueGenericType;

		/** @param keysCanBeNull False if all keys are not null. This saves 1 byte per key if keyClass is set. True if it is not known
		 *           (default). */
		public void setKeysCanBeNull (boolean keysCanBeNull) {
			this.keysCanBeNull = keysCanBeNull;
		}

		/** @param keyClass The concrete class of each key. This saves 1 byte per key. Set to null if the class is not known or varies
		 *           per key (default).
		 * @param keySerializer The serializer to use for each key. */
		public void setKeyClass (Class keyClass, Serializer keySerializer) {
			this.keyClass = keyClass;
			this.keySerializer = keySerializer;
		}

		/** @param valueClass The concrete class of each value. This saves 1 byte per value. Set to null if the class is not known or
		 *           varies per value (default).
		 * @param valueSerializer The serializer to use for each value. */
		public void setValueClass (Class valueClass, Serializer valueSerializer) {
			this.valueClass = valueClass;
			this.valueSerializer = valueSerializer;
		}

		/** @param valuesCanBeNull True if values are not null. This saves 1 byte per value if keyClass is set. False if it is not known
		 *           (default). */
		public void setValuesCanBeNull (boolean valuesCanBeNull) {
			this.valuesCanBeNull = valuesCanBeNull;
		}

		public void setGenerics (Kryo kryo, Class[] generics) {
			if (generics[0] != null && kryo.isFinal(generics[0])) keyGenericType = generics[0];
			if (generics[1] != null && kryo.isFinal(generics[1])) valueGenericType = generics[1];
		}

		public void write (Kryo kryo, Output output, Map map) {
			int length = map.size();
			output.writeInt(length, true);

			Serializer keySerializer = this.keySerializer;
			if (keyGenericType != null) {
				if (keySerializer == null) keySerializer = kryo.getSerializer(keyGenericType);
				keyGenericType = null;
			}
			Serializer valueSerializer = this.valueSerializer;
			if (valueGenericType != null) {
				if (valueSerializer == null) valueSerializer = kryo.getSerializer(valueGenericType);
				valueGenericType = null;
			}
			try{
			for (Iterator iter = map.entrySet().iterator(); iter.hasNext();) {
				Entry entry = (Entry)iter.next();
				if (keySerializer != null) {
					if (keysCanBeNull)
						kryo.writeObjectOrNull(output, ATKryo.writeReplace(entry.getKey()), keySerializer);
					else
						kryo.writeObject(output, ATKryo.writeReplace(entry.getKey()), keySerializer);
				} else
					kryo.writeClassAndObject(output, ATKryo.writeReplace(entry.getKey()));
				if (valueSerializer != null) {
					if (valuesCanBeNull)
						kryo.writeObjectOrNull(output, ATKryo.writeReplace(entry.getValue()), valueSerializer);
					else
						kryo.writeObject(output, ATKryo.writeReplace(entry.getValue()), valueSerializer);
				} else
					kryo.writeClassAndObject(output, ATKryo.writeReplace(entry.getValue()));
			}
			}catch(ObjectStreamException e){
				e.printStackTrace();
			}
		}

		/** Used by {@link #read(Kryo, Input, Class)} to create the new object. This can be overridden to customize object creation, eg
		 * to call a constructor with arguments. The default implementation uses {@link Kryo#newInstance(Class)}. */
		protected Map create (Kryo kryo, Input input, Class<Map> type) {
			return kryo.newInstance(type);
		}

		public Map read (Kryo kryo, Input input, Class<Map> type) {
			Map map = create(kryo, input, type);
			int length = input.readInt(true);

			Class keyClass = this.keyClass;
			Class valueClass = this.valueClass;

			Serializer keySerializer = this.keySerializer;
			if (keyGenericType != null) {
				keyClass = keyGenericType;
				if (keySerializer == null) keySerializer = kryo.getSerializer(keyClass);
				keyGenericType = null;
			}
			Serializer valueSerializer = this.valueSerializer;
			if (valueGenericType != null) {
				valueClass = valueGenericType;
				if (valueSerializer == null) valueSerializer = kryo.getSerializer(valueClass);
				valueGenericType = null;
			}

			kryo.reference(map);

			for (int i = 0; i < length; i++) {
				Object key;
				if (keySerializer != null) {
					if (keysCanBeNull)
						key = kryo.readObjectOrNull(input, keyClass, keySerializer);
					else
						key = kryo.readObject(input, keyClass, keySerializer);
				} else
					key = kryo.readClassAndObject(input);
				Object value;
				if (valueSerializer != null) {
					if (valuesCanBeNull)
						value = kryo.readObjectOrNull(input, valueClass, valueSerializer);
					else
						value = kryo.readObject(input, valueClass, valueSerializer);
				} else
					value = kryo.readClassAndObject(input);
				try {
					ATKryo.readResolve(key);
					ATKryo.readResolve(value);
				} catch (ObjectStreamException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
				map.put(key, value);
			}
			return map;
		}

		protected Map createCopy (Kryo kryo, Map original) {
			return kryo.newInstance(original.getClass());
		}

		public Map copy (Kryo kryo, Map original) {
			Map copy = createCopy(kryo, original);
			for (Iterator iter = original.entrySet().iterator(); iter.hasNext();) {
				Entry entry = (Entry)iter.next();
				copy.put(kryo.copy(entry.getKey()), kryo.copy(entry.getValue()));
			}
			return copy;
		}
	}
	
	
}
